#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
resplit_processed.py
基于 processed annotations.json 按病人 ID 重新生成 split.json 并写回（默认覆盖，写备份）。
用法示例:
  python resplit_processed.py --processed_dir /path/to/aier_processed --test_ratio 0.2 --seed 42

输出:
  - 在 processed_dir 下写入 split.json（格式: {"train": [...], "test": [...] }）
  - 若存在原 split.json 会自动备份为 split.json.bak.<ts>

参数:
  --processed_dir   : processed 数据目录（包含 annotations.json）
  --test_ratio      : test 集占比 (0.0-1.0)
  --seed            : 随机种子（保证可复现）
  --delimiter       : 用于从文件名提取病人 id 的分隔符 (默认 '_'，取第一个分段)
  --min_images_per_patient: 若某病人图片数 < 该阈值, 仍作为个体整体划入（默认 1）
  --force           : 即使已经存在 split.json 也覆盖（默认 False，仍会备份原文件）
"""

import argparse
import json
import os
import random
import time
from collections import defaultdict
from pathlib import Path

def parse_args():
    p = argparse.ArgumentParser(description="基于 processed annotations.json 按病人id重新划分 train/test")
    p.add_argument("--processed_dir", "-d", default="/home/zhangpinglu/data0/gy/Dataset/aier_processed", help="processed 数据目录（包含 annotations.json）")
    p.add_argument("--test_ratio", type=float, default=0.2, help="test 集占比 (0.0-1.0)")
    p.add_argument("--seed", type=int, default=42, help="随机种子，保证复现")
    p.add_argument("--delimiter", type=str, default="_", help="文件名分隔符（取第一个分段作为 patient id）")
    p.add_argument("--min_images_per_patient", type=int, default=1, help="病人下图片最小数（默认1）")
    p.add_argument("--force", action="store_true", help="强制覆盖已有 split.json（仍会备份）")
    return p.parse_args()

def load_annotations(ann_path: Path):
    if not ann_path.exists():
        raise FileNotFoundError(f"找不到文件: {ann_path}")
    with ann_path.open("r", encoding="utf-8") as f:
        data = json.load(f)
    if not isinstance(data, dict):
        raise ValueError("annotations.json 内容不是 dict")
    return data

def extract_patient_id(filename: str, delimiter: str):
    # 取第一个 delimiter 前的一段作为 patient id；若无 delimiter，则整个 filename（去扩展名）
    base = os.path.splitext(filename)[0]
    if delimiter in base:
        return base.split(delimiter, 1)[0]
    return base

def backup_file(path: Path):
    if path.exists():
        ts = int(time.time())
        bak = path.with_name(path.name + f".bak.{ts}")
        path.replace(bak)  # 原子移动
        return bak
    return None

def main():
    args = parse_args()
    processed_dir = Path(args.processed_dir)
    ann_path = processed_dir / "annotations.json"
    split_path = processed_dir / "split.json"

    data = load_annotations(ann_path)
    # 映射 patient_id -> list of image names
    pid2imgs = defaultdict(list)
    for img_name in data.keys():
        pid = extract_patient_id(img_name, args.delimiter)
        pid2imgs[pid].append(img_name)

    # 过滤极端：如果某些 pid 的图片数量 < min_images_per_patient 仍保留（默认1）
    all_pids = list(pid2imgs.keys())
    random.seed(args.seed)
    random.shuffle(all_pids)

    n_pids = len(all_pids)
    n_test_pids = max(1, int(n_pids * args.test_ratio)) if n_pids>0 else 0

    test_pids = set(all_pids[:n_test_pids])
    train_pids = set(all_pids[n_test_pids:])

    train_imgs = []
    test_imgs = []
    for pid, imgs in pid2imgs.items():
        if pid in test_pids:
            test_imgs.extend(imgs)
        else:
            train_imgs.extend(imgs)

    # 统计信息
    print(f"[info] processed_dir: {processed_dir}")
    print(f"[info] total patient ids: {n_pids}")
    print(f"[info] using seed={args.seed}, test_ratio={args.test_ratio} => test_pids={len(test_pids)}")
    print(f"[info] train images: {len(train_imgs)}, test images: {len(test_imgs)}")

    # 备份旧的 split.json（如果存在）
    if split_path.exists():
        bak = backup_file(split_path)
        print(f"[info] 备份旧 split.json -> {bak}")

    # 写入新的 split.json
    split_obj = {"train": sorted(train_imgs), "test": sorted(test_imgs)}
    with split_path.open("w", encoding="utf-8") as f:
        json.dump(split_obj, f, ensure_ascii=False, indent=2)

    print(f"[ok] 已生成 split.json -> {split_path}")

if __name__ == "__main__":
    main()
